/**
 * main.cpp
 * this file contains the implementation of calibrate onnx's from datatype of fp32 to int8.
 * 
 * usage:
 *  ./run_quantify  -input= -data= /path/to/data/dir
 * 
 * */
#include <cstdlib>
#include <cstring>
#include <iostream>
#include <fstream>
#include <cuda_runtime_api.h>
#include <NvOnnxParser.h>
#include "NvInfer.h"

#include "argsParser.h"
#include "logging.h"
#include "calibrator.h"

#define MAXBATCHSIZE 100

void print_info(){
    std::cout << "Usage: ./run_quantification  --batch [1] --useDLaCore [-1] --input_h [384] --input_w [640]"
                "--onnxfile [*.onnx] --outfile [out.engine] --datadir [] --calib_table [] --input_blob [images]"
              << std::endl;
}

void print_args(calibrate_tool::Args& args){
    std::cout << "=================args==================\n";
    std::cout << "batch : " << args.batch << '\n'
              << "useint8 : " << args.int8 << '\n'
              << "input_h : " << args.input_h << "\tinput_w : " << args.input_w << '\n'
              << "onnx_model : " << args.onnxFile << '\n'
              << "out_engine : " << args.outFile << '\n'
              << "calibrate_data_dir : " << args.dataDir << '\n'
              << "generated_calibrate_table : " << args.calib_table_name << '\n'
              << "input_blob_name : " << args.input_blob_name <<"\n";
}

int main(int argc, char* argv[])
{
    if (argc >=2 && strcmp(argv[1], "-h")==0) {print_info(); return 0;}
    calibrate_tool::Args args;
    calibrate_tool::parseArgs(args, argc, argv);
    print_args(args);
    Logger gLogger;
    // 1. build model
    nvinfer1::IBuilder* builder = nvinfer1::createInferBuilder(gLogger);
    nvinfer1::IBuilderConfig* config = builder->createBuilderConfig();
    const auto explicitBatch = 1U << static_cast<uint32_t>(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
    nvinfer1::INetworkDefinition* network = builder->createNetworkV2(explicitBatch);
    nvonnxparser::IParser* parser = nvonnxparser::createParser(*network, gLogger.getTRTLogger());
    bool parsed = parser->parseFromFile(args.onnxFile.c_str(), static_cast<int>(gLogger.getReportableSeverity()));
    if (!parsed){
        std::cout << "can't parse onnx file.\n";
        return -1;
    }
    config->setMaxWorkspaceSize(16 * (1 << 20));
    builder->setMaxBatchSize(MAXBATCHSIZE);
    nvinfer1::Dims in_dims = network->getInput(0)->getDimensions();
    bool is_fixed_shape{true};
    for(int i =0; i < in_dims.nbDims;i++){
        is_fixed_shape &= (in_dims.d[i] >=0);}
    if( !is_fixed_shape){
        auto profile = builder->createOptimizationProfile();
        profile->setDimensions(args.input_blob_name.c_str(),nvinfer1::OptProfileSelector::kMIN, nvinfer1::Dims4{1, 3, args.input_h, args.input_w});
        profile->setDimensions(args.input_blob_name.c_str(),nvinfer1::OptProfileSelector::kOPT, nvinfer1::Dims4{1, 3, args.input_h, args.input_w});
        profile->setDimensions(args.input_blob_name.c_str(),nvinfer1::OptProfileSelector::kMAX, nvinfer1::Dims4{MAXBATCHSIZE,3, args.input_h, args.input_w});
        config->addOptimizationProfile(profile);
    }
   
    if (args.int8 ==1){
        if(!builder->platformHasFastInt8()){
            std::cout << "current paltform doesn't support int8. \n";
            return -1;
        }
        // if (args.useDLACore && builder->getMaxDLABatchSize())
        config->setFlag(nvinfer1::BuilderFlag::kINT8);
        Int8EntropyCalibrator2 *calibrator = new Int8EntropyCalibrator2(args.batch, args.input_w, args.input_h, args.dataDir.c_str(), args.calib_table_name.c_str(), args.input_blob_name.c_str());
        config->setInt8Calibrator(calibrator);
        std::cout << "Building engine, please wait for a while ..." << std::endl;
    }
    nvinfer1::ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);
    if(engine)
        std::cout << "Build engine successfully!" << std::endl;
    else{
        std::cout <<"Failed to build engine.\n" << std::endl;
        return -1;
    }
    nvinfer1::IHostMemory* engine_stream{nullptr};
    engine_stream = engine->serialize();
    network->destroy();
    engine->destroy();

    assert(engine_stream != nullptr);
    std::ofstream p(args.outFile,std::ios::binary);
    if(!p){
        std::cerr << "could not open plan output file" << std::endl;
        return -1;
    }
    p.write(reinterpret_cast<const char*>(engine_stream->data()), engine_stream->size());
    engine_stream->destroy();
    return 0;
}